# width, height, top, left

import pytesseract

import numpy as np
import cv2
import sys
import time

if sys.platform == "win32":
    pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'


def formatImageOCR(originalScreenshot):
    """
    格式化图像以优化OCR(光学字符识别)。

    参数:
    originalScreenshot: 原始屏幕截图，将对此图像进行处理以增强可读性。

    返回值:
    final_image: 处理后的图像，优化了字符识别。
    """
    # 将截图转换为NumPy数组并进行预处理以找到局部最大值
    screenshot = np.array(originalScreenshot, dtype=np.uint8)
    kernelSize = 5
    maxKernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernelSize, kernelSize))
    localMax = cv2.morphologyEx(screenshot, cv2.MORPH_CLOSE, maxKernel, None, None, 1, cv2.BORDER_REFLECT101)

    # 执行增益分割以增强图像对比度
    gainDivision = np.where(localMax == 0, 0, (screenshot / localMax))
    gainDivision = np.clip((255 * gainDivision), 0, 255)
    gainDivision = gainDivision.astype("uint8")

    # 将图像从RGB转换为灰度，并调整大小以提高质量
    grayscaleImage = cv2.cvtColor(gainDivision, cv2.COLOR_BGR2GRAY)
    grayscaleImage = cv2.resize(grayscaleImage, (0, 0), fx=3.0, fy=3.0)

    # 通过Otsu方法获取二进制图像
    _, final_image = cv2.threshold(grayscaleImage, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # 执行形态学操作以关闭图像中的孔洞
    kernelSize = 3
    opIterations = 1
    morphKernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernelSize, kernelSize))
    final_image = cv2.morphologyEx(final_image, cv2.MORPH_CLOSE, morphKernel, None, None, opIterations, cv2.BORDER_REFLECT101)

    # 使用洪水填充来隔离字符区域
    cv2.floodFill(final_image, mask=None, seedPoint=(int(0), int(0)), newVal=(255))
    final_image = 255 - final_image

    # 查找并处理二值图像上的轮廓以填充字符内部的孔洞
    contours, hierarchy = cv2.findContours(final_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    for i, c in enumerate(contours):
        currentHierarchy = hierarchy[0][i][3]
        if currentHierarchy != -1:
            boundRect = cv2.boundingRect(c)
            rectX, rectY, rectWidth, rectHeight = boundRect
            fx = rectX + 0.5 * rectWidth
            fy = rectY + 0.5 * rectHeight
            cv2.floodFill(final_image, mask=None, seedPoint=(int(fx), int(fy)), newVal=(0))

    return final_image


# Change to https://stackoverflow.com/questions/66334737/pytesseract-is-very-slow-for-real-time-ocr-any-way-to-optimise-my-code 
# or https://www.reddit.com/r/learnpython/comments/kt5zzw/how_to_speed_up_pytesseract_ocr_processing/

def getTextFromImage(image):
    """
    从图像中提取文本。
    
    参数:
    image - 输入的图像，预期为可以被处理为OCR（光学字符识别）的格式。
    
    返回值:
    一个元组，包含两部分：提取到的文本字符串和处理后的图像。
    """
    imageCandidate = formatImageOCR(image)  # 格式化图像以供OCR处理
    
    # 将结果写入磁盘：
    
    # DEBUG：将当前轮次的结果写入磁盘
    # import time
    # cv2.imwrite(f"./DEBUG/{str(time.time())}.png", imageCandidate, [cv2.IMWRITE_PNG_COMPRESSION, 0])

    # 注意：这部分代码似乎存在bug
    # 使用tesseract从屏幕截图中获取当前轮次的文本
    return pytesseract.image_to_string(imageCandidate,  config='--psm 7').replace("\n", ""), imageCandidate

